import os
import cv2
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import pickle

class EmbroideryModel:
    def __init__(self):
        self.stitch_types = ["十字挑", "斜挑", "平挑", "横挑", "竖挑"]
        self.stitch_colors = {
            "十字挑": (255, 0, 0),    # 红色
            "斜挑": (0, 255, 0),     # 绿色
            "平挑": (0, 0, 255),     # 蓝色
            "横挑": (255, 255, 0),   # 黄色
            "竖挑": (255, 0, 255)    # 紫色
        }
    
    def train_model(self, train_dir):
        """
        训练随机森林分类器模型
        
        Args:
            train_dir: 训练数据目录，包含图像和标注
            
        Returns:
            训练好的模型
        """
        # 收集训练数据
        X = []
        y = []
        
        # 这里应该从训练目录中加载图像和标注
        # 简化起见，我们假设目录结构为：
        # train_dir/
        #   - stitch_type1/
        #     - image1.jpg
        #     - image2.jpg
        #   - stitch_type2/
        #     - ...
        
        for stitch_type in self.stitch_types:
            stitch_dir = os.path.join(train_dir, stitch_type)
            if not os.path.exists(stitch_dir):
                continue
                
            for img_file in os.listdir(stitch_dir):
                if not img_file.lower().endswith(('.png', '.jpg', '.jpeg')):
                    continue
                    
                img_path = os.path.join(stitch_dir, img_file)
                try:
                    # 加载图像
                    image = cv2.imread(img_path)
                    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                    
                    # 提取特征
                    from feature_extraction import FeatureExtractor
                    extractor = FeatureExtractor()
                    features = extractor.extract_features(image)
                    
                    # 添加到训练数据
                    X.append(features)
                    y.append(stitch_type)
                except Exception as e:
                    print(f"处理图像 {img_path} 时出错: {str(e)}")
        
        if len(X) == 0:
            raise ValueError("没有找到有效的训练数据")
        
        # 转换为numpy数组
        X = np.array(X)
        y = np.array(y)
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        
        # 训练随机森林分类器
        model = RandomForestClassifier(n_estimators=100, random_state=42)
        model.fit(X_train, y_train)
        
        # 评估模型
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        print(f"模型准确率: {accuracy:.4f}")
        
        return model
    
    def apply_model_prediction(self, image, grid_features, model, grid_size):
        """
        应用模型预测结果生成刺绣图
        
        Args:
            image: 原始图像
            grid_features: 网格特征字典
            model: 训练好的模型
            grid_size: 网格大小
            
        Returns:
            刺绣图像
        """
        h, w = image.shape[:2]
        result = np.zeros((h, w, 3), dtype=np.uint8)
        result.fill(255)  # 白色背景
        
        # 对每个网格进行预测并应用针法
        for (row, col), features in grid_features.items():
            # 预测针法
            stitch_type = model.predict([features])[0]
            
            # 获取网格位置
            y = row * grid_size
            x = col * grid_size
            end_y = min(y + grid_size, h)
            end_x = min(x + grid_size, w)
            
            # 应用针法
            self._draw_stitch(result, x, y, end_x, end_y, stitch_type)
        
        # 添加网格线
        for x in range(0, w, grid_size):
            cv2.line(result, (x, 0), (x, h), (200, 200, 200), 1)
        
        for y in range(0, h, grid_size):
            cv2.line(result, (0, y), (w, y), (200, 200, 200), 1)
        
        return result
    
    def apply_stitch(self, image, stitch_type, grid_size):
        """
        直接应用指定针法生成刺绣图
        
        Args:
            image: 原始图像
            stitch_type: 针法类型
            grid_size: 网格大小
            
        Returns:
            刺绣图像
        """
        h, w = image.shape[:2]
        result = np.zeros((h, w, 3), dtype=np.uint8)
        result.fill(255)  # 白色背景
        
        # 对每个网格应用针法
        for y in range(0, h, grid_size):
            for x in range(0, w, grid_size):
                end_y = min(y + grid_size, h)
                end_x = min(x + grid_size, w)
                
                # 应用针法
                self._draw_stitch(result, x, y, end_x, end_y, stitch_type)
        
        # 添加网格线
        for x in range(0, w, grid_size):
            cv2.line(result, (x, 0), (x, h), (200, 200, 200), 1)
        
        for y in range(0, h, grid_size):
            cv2.line(result, (0, y), (w, y), (200, 200, 200), 1)
        
        return result
    
    def _draw_stitch(self, image, x1, y1, x2, y2, stitch_type):
        """
        在指定区域绘制针法
        
        Args:
            image: 目标图像
            x1, y1: 左上角坐标
            x2, y2: 右下角坐标
            stitch_type: 针法类型
        """
        color = self.stitch_colors.get(stitch_type, (0, 0, 0))
        thickness = 2
        
        # 计算网格中心点和尺寸
        center_x = (x1 + x2) // 2
        center_y = (y1 + y2) // 2
        width = x2 - x1
        height = y2 - y1
        
        # 根据不同针法绘制图案
        if stitch_type == "十字挑":
            # 绘制十字
            cv2.line(image, (x1, y1), (x2, y2), color, thickness)
            cv2.line(image, (x2, y1), (x1, y2), color, thickness)
            
        elif stitch_type == "斜挑":
            # 绘制斜线
            padding = min(width, height) // 4
            cv2.line(image, (x1 + padding, y1 + padding), 
                     (x2 - padding, y2 - padding), color, thickness)
            
        elif stitch_type == "平挑":
            # 绘制垂直线
            third_w = width // 3
            for i in range(3):
                x = x1 + third_w * i + third_w // 2
                cv2.line(image, (x, y1 + height//4), (x, y2 - height//4), color, thickness)
            
        elif stitch_type == "横挑":
            # 绘制水平线
            third_h = height // 3
            for i in range(3):
                y = y1 + third_h * i + third_h // 2
                cv2.line(image, (x1 + width//4, y), (x2 - width//4, y), color, thickness)
            
        elif stitch_type == "竖挑":
            # 绘制垂直线
            third_w = width // 3
            for i in range(3):
                x = x1 + third_w * i + third_w // 2
                cv2.line(image, (x, y1 + height//4), (x, y2 - height//4), color, thickness)
        
        # 如果没有匹配的针法，绘制一个点
        else:
            cv2.circle(image, (center_x, center_y), 3, color, -1)